package com.feidee.fdspark.transformer

import org.apache.spark.ml.Transformer
import org.apache.spark.ml.param.{ParamMap, StringArrayParam}
import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable}
import org.apache.spark.sql.types.{StructField, StructType}
import org.apache.spark.sql.{DataFrame, Dataset}

/**
  * 重命名 DataFrame 指定列列名(e.g. cols=("a","b"), renameTo=("A","B"): a -> A, b -> B)
  * @Author songhaicheng
  * @Date 2018/09/11
  * @Email: haicheng_song@sui.com
  */
class ColNamer(override val uid: String) extends Transformer with DefaultParamsWritable {

  final val cols = new StringArrayParam(this, "cols", "names which will be modified")
  final val renameTo = new StringArrayParam(this, "renameTo", "names which will be display")

  def setCols(value: Array[String]): this.type = set(cols, value)
  def getCols: Array[String] = $(cols)
  def setRenameTo(value: Array[String]): this.type = set(renameTo, value)
  def getRenameTo: Array[String] = $(renameTo)

  // 默认为空
  setDefault(cols, Array[String]())
  setDefault(renameTo, Array[String]())

  def this() = this(Identifiable.randomUID("colnamer"))

  override def transform(dataset: Dataset[_]): DataFrame = {
    val columns = dataset.schema.fieldNames.map(name => {
      val colIdx = $(cols).indexOf(name)
      if (colIdx != -1) {
        dataset.col(name).as($(renameTo)(colIdx))
      } else {
        dataset.col(name)
      }
    })
    dataset.select(columns:_*)
  }

  override def copy(extra: ParamMap): Transformer = {
    defaultCopy(extra)
  }

  override def transformSchema(schema: StructType): StructType = {
    val names = schema.fieldNames
    // 参数不能为空
    require($(cols).length > 0 && $(renameTo).length > 0, "params must be all set")
    // 被重命名列和重命名列必须一一对应
    require($(cols).length == $(renameTo).length, "cols' length must be equals to renameTo's length")
    // cols 必须存在
    $(cols).foreach(col => require(names.contains(col), s"item of cols must be existed (false item: $col)"))

    // 进行重命名
    val columns = schema.fields.map(field => {
      val colIdx = $(cols).indexOf(field.name)
      if (colIdx != -1) {
        StructField($(renameTo)(colIdx), field.dataType, field.nullable)
      } else {
        field
      }
    })
    StructType(columns)
  }

}

object ColNamer extends DefaultParamsReadable[ColNamer] {}
